from typing import List, Tuple

import torch
from wenet.utils.common import log_add


class Sequence():

    __slots__ = {'hyp', 'score', 'cache'}

    def __init__(
        self,
        hyp: List[torch.Tensor],
        score,
        cache: List[torch.Tensor],
    ):
        self.hyp = hyp
        self.score = score
        self.cache = cache


def shellSortHyps(nums: List[Sequence]) -> List[Sequence]:
    lens = len(nums)
    gap: int = 1
    while gap < lens // 3:
        gap = gap * 3 + 1  # 动态定义间隔序列
    while gap > 0:
        for i in range(gap, lens):
            # curNum storing the current inserting number
            curNum, preIndex = nums[i].score, i - gap
            while preIndex >= 0 and curNum > nums[preIndex].score:
                # make less than curNum element to the last index
                nums[preIndex + gap].score = nums[preIndex].score
                preIndex -= gap
            # correct index for inserting number
            nums[preIndex + gap].score = curNum
        gap = gap // 3  # next dynamic interval
    return nums


class PrefixBeamSearchStreaming():

    def __init__(self, encoder, predictor, joint, ctc, blank):
        self.encoder = encoder
        self.predictor = predictor
        self.joint = joint
        self.ctc = ctc
        self.blank = blank
        self.device = torch.device("cpu")
        self.beam_init: List[Sequence] = []
        self.cache = self.predictor.init_state(
            1, method="zero", device=self.device)
        self.beam_init.append(
            Sequence(hyp=[self.blank], score=0.0, cache=self.cache))

        self.att_cache: torch.Tensor = torch.zeros(
            (0, 0, 0, 0), device=self.device)
        self.cnn_cache: torch.Tensor = torch.zeros(
            (0, 0, 0, 0), device=self.device)
        self.feat_offset: int = 0

    def forward_decoder_one_step(
            self, encoder_x: torch.Tensor, pre_t: torch.Tensor,
            cache: List[torch.Tensor]
    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        padding = torch.zeros(pre_t.size(0), 1, device=encoder_x.device)
        pre_t, new_cache = self.predictor.forward_step(pre_t.unsqueeze(-1),
                                                       padding, cache)
        x = self.joint(encoder_x, pre_t)  # [beam, 1, 1, vocab]
        x = x.log_softmax(dim=-1)
        return x, new_cache

    def prefix_beam_search(self,
                           speech: torch.Tensor,
                           speech_lengths: torch.Tensor,
                           decoding_chunk_size: int = 16,
                           beam_size: int = 1,
                           num_decoding_left_chunks: int = -1,
                           last_chunk: bool = False,
                           ctc_weight: float = 0.3,
                           transducer_weight: float = 0.7):
        """prefix beam search
           also see wenet.transducer.transducer.beam_search
        """
        assert speech.shape[0] == speech_lengths.shape[0]
        assert decoding_chunk_size != 0
        device = speech.device
        batch_size = speech.shape[0]
        assert batch_size == 1

        # 1. Encoder
        # encoder_out, _ = self.encoder(
        #     speech, speech_lengths, decoding_chunk_size,
        #     num_decoding_left_chunks)  # (B, maxlen, encoder_dim)

        encoder_out, _ = self.encoder.forward_chunk_by_chunk_simple(
            speech, decoding_chunk_size,
            num_decoding_left_chunks)  # (B, maxlen, encoder_dim)

        maxlen = encoder_out.size(1)

        ctc_probs = self.ctc.log_softmax(encoder_out).squeeze(0)
        # beam_init: List[Sequence] = []

        # 2. init beam using Sequence to save beam unit
        # cache = self.predictor.init_state(1, method="zero", device=device)
        self.beam_init.append(
            Sequence(hyp=[self.blank], score=0.0, cache=self.cache))
        # 3. start decoding (notice: we use breathwise first searching)
        # !!!! In this decoding method: one frame do not output multi units. !!!!
        # !!!!    Experiments show that this strategy has little impact      !!!!
        for i in range(maxlen):
            # 3.1 building input
            # decoder taking the last token to predict the next token
            input_hyp = [s.hyp[-1] for s in self.beam_init]
            input_hyp_tensor = torch.tensor(input_hyp,
                                            dtype=torch.int,
                                            device=device)
            # building statement from beam. return a List of tensors !!!!
            cache_batch = self.predictor.cache_to_batch(
                [s.cache for s in self.beam_init])
            # build score tensor to do torch.add() function
            scores = torch.tensor([s.score for s in self.beam_init]).to(device)

            # 3.2 forward decoder
            logp, new_cache = self.forward_decoder_one_step(
                encoder_out[:, i, :].unsqueeze(1),
                input_hyp_tensor,
                cache_batch,
            )  # logp: (N, 1, 1, vocab_size)
            logp = logp.squeeze(1).squeeze(1)  # logp: (N, vocab_size)
            new_cache = self.predictor.batch_to_cache(new_cache)

            # 3.3 shallow fusion for transducer score
            #     and ctc score where we can also add the LM score
            # transducer_weight == 0.0, so logp = ctc_probs
            logp = torch.log(
                torch.add(transducer_weight * torch.exp(logp),
                          ctc_weight * torch.exp(ctc_probs[i].unsqueeze(0))))

            # 3.4 first beam prune
            top_k_logp, top_k_index = logp.topk(beam_size)  # (N, N)
            scores = torch.add(scores.unsqueeze(1), top_k_logp)

            # 3.5 generate new beam (N*N)
            beam_A = []
            for j in range(len(self.beam_init)):
                # update seq
                base_seq = self.beam_init[j]
                for t in range(beam_size):
                    # blank: only update the score
                    if top_k_index[j, t] == self.blank:
                        new_seq = Sequence(hyp=base_seq.hyp.copy(),
                                           score=scores[j, t].item(),
                                           cache=base_seq.cache)

                        beam_A.append(new_seq)
                    # other unit: update hyp score statement and last
                    else:
                        hyp_new = base_seq.hyp.copy()
                        hyp_new.append(top_k_index[j, t].item())
                        new_seq = Sequence(hyp=hyp_new,
                                           score=scores[j, t].item(),
                                           cache=new_cache[j])
                        beam_A.append(new_seq)

            # 3.6 prefix fusion
            fusion_A = [beam_A[0]]
            for j in range(1, len(beam_A)):
                s1 = beam_A[j]
                if_do_append = True
                for t in range(len(fusion_A)):
                    # notice: A_ can not fusion with A
                    if s1.hyp == fusion_A[t].hyp:
                        fusion_A[t].score = log_add(
                            [fusion_A[t].score, s1.score])
                        if_do_append = False
                        break
                if if_do_append:
                    fusion_A.append(s1)

            # 4. second pruned
            fusion_A.sort(key=lambda x: x.score, reverse=True)
            self.beam_init = fusion_A[:beam_size]
        self.cache = new_cache
        return self.beam_init, encoder_out

    @torch.jit.export
    def prefix_beam_search_chunk(self,
                                 speech: torch.Tensor,
                                 speech_lengths: torch.Tensor,
                                 decoding_chunk_size: int = 16,
                                 beam_size: int = 1,
                                 num_decoding_left_chunks: int = -1,
                                 last_chunk: bool = False,
                                 ctc_weight: float = 0.3,
                                 transducer_weight: float = 0.7):
        """prefix beam search
           also see wenet.transducer.transducer.beam_search
        """
        assert speech.shape[0] == speech_lengths.shape[0]
        assert decoding_chunk_size != 0
        device = speech.device
        batch_size = speech.shape[0]
        assert batch_size == 1

        # 0. key parameters defined here
        subsampling = 4
        right_context = 6
        context = right_context + 1  # Add current frame
        # key parameters defined ending
        stride = subsampling * decoding_chunk_size
        decoding_window = (decoding_chunk_size -
                           1) * subsampling + context
        num_frames = speech.size(1)
        offset = 0
        required_cache_size = decoding_chunk_size * num_decoding_left_chunks
        outputs = []
        # Feed forward overlap input step by step
        for cur in range(0, num_frames - context + 1, stride):
            end = min(cur + decoding_window, num_frames)
            chunk_xs = speech[:, cur:end, :]
            (encoder_out, self.att_cache, self.cnn_cache) = self.encoder.forward_chunk(
                chunk_xs, offset, required_cache_size, self.att_cache, self.cnn_cache)
            outputs.append(encoder_out)
            offset += encoder_out.size(1)
            segment_maxlen = encoder_out.size(1)
            ctc_probs = self.ctc.log_softmax(encoder_out).squeeze(0)
            # to beam search here
            for i in range(segment_maxlen):
                # 3.1 building input
                # decoder taking the last token to predict the next token
                input_hyp = [s.hyp[-1] for s in self.beam_init]
                input_hyp_tensor = torch.tensor(input_hyp,
                                                dtype=torch.int,
                                                device=device)
                # building statement from beam. return a List of tensors !!!!
                cache_batch = self.predictor.cache_to_batch(
                    [s.cache for s in self.beam_init])
                # build score tensor to do torch.add() function
                scores = torch.tensor(
                    [s.score for s in self.beam_init]).to(device)

                # 3.2 forward decoder
                logp, new_cache = self.forward_decoder_one_step(
                    encoder_out[:, i, :].unsqueeze(1),
                    input_hyp_tensor,
                    cache_batch,
                )  # logp: (N, 1, 1, vocab_size)
                logp = logp.squeeze(1).squeeze(1)  # logp: (N, vocab_size)
                new_cache = self.predictor.batch_to_cache(new_cache)

                # 3.3 shallow fusion for transducer score
                #     and ctc score where we can also add the LM score
                # transducer_weight == 0.0, so logp = ctc_probs
                logp = torch.log(
                    torch.add(transducer_weight * torch.exp(logp),
                              ctc_weight * torch.exp(ctc_probs[i].unsqueeze(0))))

                # 3.4 first beam prune
                top_k_logp, top_k_index = logp.topk(beam_size)  # (N, N)
                scores = torch.add(scores.unsqueeze(1), top_k_logp)

                # 3.5 generate new beam (N*N)
                # start to search for new beam
                beam_A = []
                for j in range(len(self.beam_init)):
                    # update seq
                    base_seq = self.beam_init[j]
                    for t in range(beam_size):
                        # blank: only update the score
                        if top_k_index[j, t] == self.blank:
                            new_seq = Sequence(hyp=base_seq.hyp.copy(),
                                               score=scores[j, t].item(),
                                               cache=base_seq.cache)

                            beam_A.append(new_seq)
                        # other unit: update hyp score statement and last
                        else:
                            hyp_new = base_seq.hyp.copy()
                            hyp_new.append(top_k_index[j, t].item())
                            new_seq = Sequence(hyp=hyp_new,
                                               score=scores[j, t].item(),
                                               cache=new_cache[j])
                            beam_A.append(new_seq)

                # 3.6 prefix fusion
                fusion_A = [beam_A[0]]
                for j in range(1, len(beam_A)):
                    s1 = beam_A[j]
                    if_do_append = True
                    for t in range(len(fusion_A)):
                        # notice: A_ can not fusion with A
                        if s1.hyp == fusion_A[t].hyp:
                            fusion_A[t].score = log_add(
                                [fusion_A[t].score, s1.score])
                            if_do_append = False
                            break
                    if if_do_append:
                        fusion_A.append(s1)

                # 4. second pruned
                fusion_A.sort(key=lambda x: x.score, reverse=True)
                self.beam_init = fusion_A[:beam_size]

        return self.beam_init, outputs

    @torch.jit.export
    def prefix_beam_search_chunk_stream(self,
                                        speech: torch.Tensor,
                                        decoding_chunk_size: int = 16,
                                        # subsampling: int = 4,
                                        # right_context: int = 6,
                                        beam_size: int = 1,
                                        num_decoding_left_chunks: int = -1,
                                        # last_chunk: bool = False,
                                        ctc_weight: float = 0.3,
                                        transducer_weight: float = 0.7):
        """prefix beam search
           also see wenet.transducer.transducer.beam_search
        """
        # assert speech.shape[0] == speech_lengths.shape[0]
        assert decoding_chunk_size != 0
        device = speech.device
        # batch_size = speech.shape[0]
        # assert batch_size == 1

        # 0. key parameters defined here
        # subsampling = 4
        # right_context = 6
        # context = right_context + 1  # Add current frame
        # key parameters defined ending
        # stride = subsampling * decoding_chunk_size
        # decoding_window = (decoding_chunk_size -
        #                    1) * subsampling + context
        # num_frames = speech.size(1)
        # offset = 0
        required_cache_size = decoding_chunk_size * num_decoding_left_chunks
        # outputs = []
        # Feed forward overlap input step by step
        # for cur in range(0, num_frames - context + 1, stride):
        # end = min(cur + decoding_window, num_frames)
        # chunk_xs = speech[:, cur:end, :]
        # chunk_xs = speech
        (encoder_out, self.att_cache, self.cnn_cache) = self.encoder.forward_chunk(
            speech, self.feat_offset, required_cache_size, self.att_cache, self.cnn_cache)
        # outputs.append(encoder_out)
        self.feat_offset += encoder_out.size(1)
        segment_maxlen = encoder_out.size(1)
        ctc_probs = self.ctc.log_softmax(encoder_out).squeeze(0)
        # to beam search here
        for i in range(segment_maxlen):
            # 3.1 building input
            # decoder taking the last token to predict the next token
            input_hyp = [s.hyp[-1] for s in self.beam_init]
            input_hyp_tensor = torch.tensor(input_hyp,
                                            dtype=torch.int,
                                            device=device)
            # building statement from beam. return a List of tensors !!!!
            cache_batch = self.predictor.cache_to_batch(
                [s.cache for s in self.beam_init])
            # build score tensor to do torch.add() function
            scores = torch.tensor(
                [s.score for s in self.beam_init]).to(device)

            # 3.2 forward decoder
            logp, new_cache = self.forward_decoder_one_step(
                encoder_out[:, i, :].unsqueeze(1),
                input_hyp_tensor,
                cache_batch,
            )  # logp: (N, 1, 1, vocab_size)
            logp = logp.squeeze(1).squeeze(1)  # logp: (N, vocab_size)
            new_cache = self.predictor.batch_to_cache(new_cache)

            # 3.3 shallow fusion for transducer score
            #     and ctc score where we can also add the LM score
            # transducer_weight == 0.0, so logp = ctc_probs
            logp = torch.log(
                torch.add(transducer_weight * torch.exp(logp),
                          ctc_weight * torch.exp(ctc_probs[i].unsqueeze(0))))

            # 3.4 first beam prune
            top_k_logp, top_k_index = logp.topk(beam_size)  # (N, N)
            scores = torch.add(scores.unsqueeze(1), top_k_logp)

            # 3.5 generate new beam (N*N)
            # start to search for new beam
            beam_A = []
            for j in range(len(self.beam_init)):
                # update seq
                base_seq = self.beam_init[j]
                for t in range(beam_size):
                    # blank: only update the score
                    if top_k_index[j, t] == self.blank:
                        new_seq = Sequence(hyp=base_seq.hyp.copy(),
                                           score=scores[j, t].item(),
                                           cache=base_seq.cache)

                        beam_A.append(new_seq)
                    # other unit: update hyp score statement and last
                    else:
                        hyp_new = base_seq.hyp.copy()
                        hyp_new.append(top_k_index[j, t].item())
                        new_seq = Sequence(hyp=hyp_new,
                                           score=scores[j, t].item(),
                                           cache=new_cache[j])
                        beam_A.append(new_seq)

            # 3.6 prefix fusion
            fusion_A = [beam_A[0]]
            for j in range(1, len(beam_A)):
                s1 = beam_A[j]
                if_do_append = True
                for t in range(len(fusion_A)):
                    # notice: A_ can not fusion with A
                    if s1.hyp == fusion_A[t].hyp:
                        fusion_A[t].score = log_add(
                            [fusion_A[t].score, s1.score])
                        if_do_append = False
                        break
                if if_do_append:
                    fusion_A.append(s1)

            # 4. second pruned
            # raw codes here
            # fusion_A.sort(key=lambda x: x.score, reverse=True)
            # added by gaoxinglong
            sorted_fusion_A = shellSortHyps(fusion_A)
            self.beam_init = sorted_fusion_A[:beam_size]

        return self.beam_init, encoder_out

    def prefix_beam_search_chunk_by_chunk(self,
                                          speech: torch.Tensor,
                                          speech_lengths: torch.Tensor,
                                          decoding_chunk_size: int = 16,
                                          beam_size: int = 1,
                                          num_decoding_left_chunks: int = -1,
                                          last_chunk: bool = False,
                                          ctc_weight: float = 0.3,
                                          transducer_weight: float = 0.7):
        """prefix beam search
           also see wenet.transducer.transducer.beam_search
        """
        assert speech.shape[0] == speech_lengths.shape[0]
        assert decoding_chunk_size != 0
        device = speech.device
        batch_size = speech.shape[0]
        assert batch_size == 1

        # 0. some useful parameters
        subsampling = 4
        right_context = 6
        context = right_context + 1  # Add current frame
        stride = subsampling * decoding_chunk_size
        decoding_window = (decoding_chunk_size -
                           1) * subsampling + context
        num_frames = speech.size(1)
        offset = 0
        required_cache_size = decoding_chunk_size * num_decoding_left_chunks

        # 1. Encoder
        # encoder_out, _ = self.encoder(
        #     speech, speech_lengths, decoding_chunk_size,
        #     num_decoding_left_chunks)  # (B, maxlen, encoder_dim)

        encoder_out, _ = self.encoder.forward_chunk_by_chunk_simple(
            speech, decoding_chunk_size,
            num_decoding_left_chunks)  # (B, maxlen, encoder_dim)

        maxlen = encoder_out.size(1)
        # for current model, max_len is 16 for Every batch

        ctc_probs = self.ctc.log_softmax(encoder_out).squeeze(0)
        # beam_init: List[Sequence] = []

        # 2. init beam using Sequence to save beam unit
        # cache = self.predictor.init_state(1, method="zero", device=device)
        self.beam_init.append(
            Sequence(hyp=[self.blank], score=0.0, cache=self.cache))
        # 3. start decoding (notice: we use breathwise first searching)
        # !!!! In this decoding method: one frame do not output multi units. !!!!
        # !!!!    Experiments show that this strategy has little impact      !!!!
        for i in range(maxlen):
            # 3.1 building input
            # decoder taking the last token to predict the next token
            input_hyp = [s.hyp[-1] for s in self.beam_init]
            input_hyp_tensor = torch.tensor(input_hyp,
                                            dtype=torch.int,
                                            device=device)
            # building statement from beam. return a List of tensors !!!!
            cache_batch = self.predictor.cache_to_batch(
                [s.cache for s in self.beam_init])
            # build score tensor to do torch.add() function
            scores = torch.tensor([s.score for s in self.beam_init]).to(device)

            # 3.2 forward decoder
            logp, new_cache = self.forward_decoder_one_step(
                encoder_out[:, i, :].unsqueeze(1),
                input_hyp_tensor,
                cache_batch,
            )  # logp: (N, 1, 1, vocab_size)
            logp = logp.squeeze(1).squeeze(1)  # logp: (N, vocab_size)
            new_cache = self.predictor.batch_to_cache(new_cache)

            # 3.3 shallow fusion for transducer score
            #     and ctc score where we can also add the LM score
            # transducer_weight == 0.0, so logp = ctc_probs
            logp = torch.log(
                torch.add(transducer_weight * torch.exp(logp),
                          ctc_weight * torch.exp(ctc_probs[i].unsqueeze(0))))

            # 3.4 first beam prune
            top_k_logp, top_k_index = logp.topk(beam_size)  # (N, N)
            scores = torch.add(scores.unsqueeze(1), top_k_logp)

            # 3.5 generate new beam (N*N)
            beam_A = []
            for j in range(len(self.beam_init)):
                # update seq
                base_seq = self.beam_init[j]
                for t in range(beam_size):
                    # blank: only update the score
                    if top_k_index[j, t] == self.blank:
                        new_seq = Sequence(hyp=base_seq.hyp.copy(),
                                           score=scores[j, t].item(),
                                           cache=base_seq.cache)

                        beam_A.append(new_seq)
                    # other unit: update hyp score statement and last
                    else:
                        hyp_new = base_seq.hyp.copy()
                        hyp_new.append(top_k_index[j, t].item())
                        new_seq = Sequence(hyp=hyp_new,
                                           score=scores[j, t].item(),
                                           cache=new_cache[j])
                        beam_A.append(new_seq)

            # 3.6 prefix fusion
            fusion_A = [beam_A[0]]
            for j in range(1, len(beam_A)):
                s1 = beam_A[j]
                if_do_append = True
                for t in range(len(fusion_A)):
                    # notice: A_ can not fusion with A
                    if s1.hyp == fusion_A[t].hyp:
                        fusion_A[t].score = log_add(
                            [fusion_A[t].score, s1.score])
                        if_do_append = False
                        break
                if if_do_append:
                    fusion_A.append(s1)

            # 4. second pruned
            fusion_A.sort(key=lambda x: x.score, reverse=True)
            self.beam_init = fusion_A[:beam_size]
        self.cache = new_cache
        return self.beam_init, encoder_out
